import json
from elasticsearch import Elasticsearch, exceptions
from datetime import datetime
import pandas as pd

# 多集群配置 (支持密码认证)
CLUSTERS = [
    {
        "name": "万达集群",
        "url": "http://172.16.100.76:9200",
        "auth": ("elastic", "Elastic#9200")
    }
]
# CLUSTERS = [
#     {
#         "name": "生产集群",
#         "url": "http://10.2.158.45:9200",
#         "auth": None  # 无认证
#     },
#     {
#         "name": "APM集群",
#         "url": "http://10.2.158.105:9200",
#         "auth": ("elastic", "Elastic#9200")
#     }
# ]


def get_es_client(cluster_cfg):
    """创建ES客户端连接"""
    try:
        return Elasticsearch(
            hosts=[cluster_cfg["url"]],
            http_auth=cluster_cfg.get("auth"),
            verify_certs=False,
            request_timeout=30
        )
    except Exception as e:
        print(f"连接集群 {cluster_cfg['name']} 失败: {str(e)}")
        return None


def check_cluster_health(es):
    """检查集群健康状态"""
    try:
        health = es.cluster.health()
        return {
            "status": health["status"],
            "node_count": health["number_of_nodes"],
            "index_count": health["number_of_indices"]
        }
    except exceptions.AuthorizationException as e:
        return {"error": str(e)}


def get_indices_status(es):
    """获取所有索引状态"""
    try:
        indices = es.cat.indices(format="json", h="index,health")
        non_green = [idx["index"] for idx in indices if idx["health"] != "green"]
        return non_green
    except Exception as e:
        return []


def get_nodes_stats(es):
    """获取节点资源统计"""
    try:
        stats = es.nodes.stats(metric="os,jvm,fs")
        nodes = stats["nodes"].values()

        cpu_usage = max([n["os"]["cpu"]["percent"] for n in nodes])
        mem_usage = max([n["jvm"]["mem"]["heap_used_percent"] for n in nodes])

        total_storage = sum([n["fs"]["total"]["total_in_bytes"] for n in nodes])
        used_storage = sum([n["fs"]["total"]["used_in_bytes"] for n in nodes])

        return {
            "max_cpu": cpu_usage,
            "max_mem": mem_usage,
            "total_storage_gb": round(total_storage / 1024 ** 3, 2),
            "used_storage_gb": round(used_storage / 1024 ** 3, 2)
        }
    except Exception:
        return {}


def get_users(es):
    """获取用户列表"""
    try:
        users = es.security.get_user()
        return list(users.keys())
    except Exception:
        return ["安全功能未启用"]


def collect_cluster_data(cluster_cfg):
    """收集单个集群数据"""
    es = get_es_client(cluster_cfg)
    if not es:
        return None

    data = {"cluster_name": cluster_cfg["name"]}

    # 基础信息
    data.update(check_cluster_health(es))

    # 资源使用
    stats = get_nodes_stats(es)
    data.update({
        "cpu_usage(%)": stats.get("max_cpu", 0),
        "mem_usage(%)": stats.get("max_mem", 0),
        "storage_usage": f"{stats.get('used_storage_gb', 0)}GB / {stats.get('total_storage_gb', 0)}GB",
        "storage_percent(%)": round(
            (stats.get("used_storage_gb", 0) / stats.get("total_storage_gb", 1)) * 100, 2
        ) if stats.get("total_storage_gb", 0) > 0 else 0
    })

    # 索引状态
    non_green = get_indices_status(es)
    data["non_green_indices"] = "\n".join(non_green) if non_green else "无"
    data["non_green_count"] = len(non_green)

    # 用户信息
    data["users"] = "\n".join(get_users(es))

    # 版本信息
    try:
        data["version"] = es.info()["version"]["number"]
    except:
        data["version"] = "未知"

    return data


# def generate_report():
#     """生成巡检报告"""
#     report_data = []
#
#     for cluster in CLUSTERS:
#         print(f"正在检查集群: {cluster['name']}")
#         cluster_data = collect_cluster_data(cluster)
#         if cluster_data:
#             report_data.append(cluster_data)
#
#     # 创建DataFrame
#     df = pd.DataFrame(report_data)
#
#     # 列顺序调整
#     columns = [
#         "cluster_name", "version", "status", "node_count", "index_count",
#         "cpu_usage(%)", "mem_usage(%)", "storage_usage", "storage_percent(%)",
#         "non_green_count", "non_green_indices", "users"
#     ]
#     df = df[columns]
#
#     # 生成Excel文件名
#     timestamp = datetime.now().strftime("%Y%m%d_%H%M")
#     filename = f"ES_巡检报告_{timestamp}.xlsx"
#
#     # 写入Excel
#     with pd.ExcelWriter(filename, engine="xlsxwriter") as writer:
#         df.to_excel(writer, sheet_name="集群概览", index=False)
#
#         # 添加格式
#         workbook = writer.book
#         worksheet = writer.sheets["集群概览"]
#
#         # 设置列宽
#         for idx, col in enumerate(df.columns):
#             width = max(len(str(col)) + 2, 15)
#             worksheet.set_column(idx, idx, width)
#
#         # 状态列颜色标记
#         format_red = workbook.add_format({"bg_color": "#FFC7CE"})
#         format_green = workbook.add_format({"bg_color": "#C6EFCE"})
#
#         for row_num in range(1, len(df) + 1):
#             cell_value = df.iloc[row_num - 1]["status"]
#             fmt = format_green if cell_value == "green" else format_red
#             worksheet.conditional_format(
#                 row_num, 3, row_num, 3,
#                 {"type": "no_errors", "format": fmt}
#             )
#
#     print(f"报告已生成: {filename}")


if __name__ == "__main__":
    #generate_report()
    for cluster in CLUSTERS:
        print(f"正在检查集群: {cluster['name']}")
        cluster_data = collect_cluster_data(cluster)
        print(cluster_data)